import collections
import pandas as pd
import numpy as np


def get_decoder(method):
    if method == 'bio':
        return decode_bio
    if method == 'bios':
        return decode_bios
    if method == 'span':
        return decode_span


class SpanF1:
    def __init__(self, method='bio', id2label=None):
        self.metrics = collections.defaultdict(lambda: [0, 0, 0])
        self.id2label = id2label
        self.decoder = get_decoder(method)

    def update(self, yt, yp, batch):
        yt = yt.cpu().numpy()
        yp = yp.cpu().numpy()
        seq_len = batch['attention_mask'].sum(1).cpu().numpy()
        for t, p, l in zip(yt, yp, seq_len):
            t = set(self.decoder(t, l))
            p = set(self.decoder(p, l))
            tp_pair = p & t

            self.metrics['gross'][0] += len(tp_pair)
            self.metrics['gross'][1] += len(t)
            self.metrics['gross'][2] += len(p)
            for i in tp_pair:
                self.metrics[i[0]][0] += 1
            for i in t:
                self.metrics[i[0]][1] += 1
            for i in p:
                self.metrics[i[0]][2] += 1

    def accumulate(self):
        tmp = []
        index = []
        for k, (tp, num_t, num_p) in self.metrics.items():
            precision = tp / num_p if num_p > 0 else 0
            recall = tp / num_t if num_t > 0 else 0
            f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0
            tmp.append({'precision': round(precision, 4), 'recall': round(recall, 4), 'f1': round(f1, 4),
                        'support': num_t})
            index.append(k if (self.id2label is None or k == 'gross') else self.id2label[k])
            if k == 'gross':
                main_metrics = f1
        self.reset()
        print(pd.DataFrame(tmp, index=index))
        return main_metrics

    def reset(self):
        self.metrics = collections.defaultdict(lambda: [0, 0, 0])


def decode_bio(label, seq_len):
    start = False
    ret = []
    for i, v in enumerate(label[:seq_len]):
        l = (v - 1) // 2
        if v % 2 == 1:
            ret.append([l, i, i])
            start = True
        elif start and v != 0 and l == ret[-1][0]:
            ret[-1][-1] = i
        else:
            start = False
    return [tuple(x) for x in ret]


def decode_bios(label, seq_len):
    start = False
    ret = []
    for i, v in enumerate(label[:seq_len]):
        l = (v - 1) // 3
        if v % 3 == 1:
            ret.append([l, i, i])
            start = True
        elif v % 3 == 2 and start and l == ret[-1][0]:
            ret[-1][-1] = i
        elif v % 3 == 0 and v != 0:
            ret.append([l, i, i])
            start = False
        else:
            start = False
    return [tuple(x) for x in ret]


def decode_span(label, seq_len):
    ret = []
    label = label[1:seq_len]
    start, end = np.split(label, 2, -1)
    if start.shape[-1] != 1:
        start = start.argmax(-1)
        end = end.argmax(-1)
    else:
        start = start.squeeze(-1)
        end = end.squeeze(-1)
    for i, s_l in enumerate(start):
        if s_l == 0:
            continue
        for j, e_l in enumerate(end[i:]):
            if s_l == e_l:
                ret.append((s_l - 1, i, i + j))
                break
    return ret


def get_result(outputs, texts, offsets, id2label):
    ret = []
    for output, text, offset in zip(outputs, texts, offsets):
        tmp = []
        for label_id, start, end in output:
            start = offset[start][0]
            end = offset[end][1]
            tmp.append([text[start: end], id2label[label_id], start, end])
        ret.append(tmp)
    return ret


#
#         # elif method == 'span':
#         #     for s, e, t, o in zip(self.start_logits, self.end_logits, self.texts, self.offset):
#         #         s = s.cpu().numpy()[1:-1]
#         #         e = e.cpu().numpy()[1:-1]
#         #         o = o[1:-1]
#         #         for i, s_l in enumerate(s):
#         #             if s_l == 0:
#         #                 continue
#         #             for j, e_l in enumerate(e[i:]):
#         #                 if s_l == e_l:
#         #                     mapping_start = o[i][0]
#         #                     mapping_end = o[i + j][1]
#         #                     entity_text = t[mapping_start:mapping_end]
#         #                     output.append({"entity": entity_text,
#         #                                    'label': self.id2label[s_l - 1],
#         #                                    'spans': [(mapping_start, mapping_end - 1)]})
#         #                     break
#
#         return output
#
#
# class BiAffineDecoder(OutputDecoder):
#     def __call__(self, logits, texts, offsets):
#         output = []
#         labels = self.generate_label(logits)
#         for text, offset, label in zip(texts, offsets, labels):
#             line = {}
#             seq_len = len(offset)
#             indices_mat = torch.triu(torch.ones((seq_len, seq_len)), diagonal=0, out=None).to(labels[0].device)
#             indices = torch.where(indices_mat > 0)
#             logits = label[indices[0], indices[1]].argmax(-1)
#
#             label_mapping = {}
#             for i in range(seq_len):
#                 for j in range(i, seq_len):
#                     label_mapping[len(label_mapping)] = (i, j)
#
#             for idx, value in enumerate(logits):
#                 if value > 0:
#                     start, end = label_mapping[idx]
#                     mapping_start = offset[start][0]
#                     mapping_end = offset[end][1]
#                     entity_text = text[mapping_start:mapping_end]
#                     if entity_text in line:
#                         line[entity_text]['spans'].append((mapping_start, mapping_end - 1))
#                     else:
#                         line[entity_text] = {"entity": entity_text, 'label': self.id2label[value - 1],
#                                              'spans': [(mapping_start, mapping_end - 1)]}
#             output.append(list(line.values()))
#         return output
#
#


if __name__ == '__main__':
    from paddlenlp.transformers import BertTokenizer

    x = [0, 0, 0, 1, 2, 3, 4, 6, 4, 5, 6, 1]
    print(decode_bio(x))

    x = [0, 1, 2, 3, 4, 5, 6, 7, 86, 9, 56, 0]
    print(decode_bios(x))
